import sys

import torch

import torch.optim as optim
from pytorch.模型定义 import SimpleNN

# 假设已经定义好了模型、损失函数和优化器
# 定义模型
model = SimpleNN()

# 定义损失函数 均方误差损失
criterion = torch.nn.MSELoss()

# 定义优化器
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练数据示例
X = torch.randn(10, 2)  # 10 个样本，每个样本有 2 个特征
print(f'X：{X}')
Y = torch.randn(10, 1)  # 10 个目标标签
print("Y：", Y)

# 只进行一次epoch
model.train()  # 设置模型为训练模式
optimizer.zero_grad()  # 清除梯度
output = model(X)  # 前向传播
print("output：", output)
loss = criterion(output, Y)  # 计算损失
print("loss：", loss)
loss.backward()  # 反向传播
optimizer.step()  # 更新权重
print(f'Epoch , Loss: {loss.item():.4f}')




sys.exit()



# 训练过程
for epoch in range(100):  # 训练 100 轮
    model.train()  # 设置模型为训练模式
    optimizer.zero_grad()  # 清除梯度
    output = model(X)  # 前向传播
    print("output：", output)
    loss = criterion(output, Y)  # 计算损失
    loss.backward()  # 反向传播
    optimizer.step()  # 更新权重

    if (epoch + 1) % 10 == 0:  # 每 10 轮输出一次损失
        print(f'Epoch [{epoch + 1}/100], Loss: {loss.item():.4f}')
